import numpy as np
from matplotlib import pyplot as plt
import random
# from FitnessPlot import FitnessPlot
from CEC.CEC_Fun import CEC
################################################################
# 鲸鱼优化算法
################################################################


class Whale():
    """鲸鱼类"""

    def __init__(self, position):
        self.position = position
        self.fitness = 0.0


class WOA():
    """鲸鱼优化算法"""

    def __init__(self, FitFunction, LB, UB, dim, population_size, b, max_iter):
        """初始化鲸鱼优化算法的参数"""
        self.LB = LB  # 搜索空间的左边界
        self.UB = UB  # 搜索空间的右边界
        self.dim = dim  # 位置信息的维度
        self.population_size = population_size  # 种群大小
        self.max_iter = max_iter  # 迭代次数
        self.b = b  # 螺旋式更新位置的参数
        self.X = np.zeros((population_size, dim))  # 种群的位置信息
        self.FitFunction = FitFunction
        self.gBest_score = np.inf  # 最优解对应的适应度值
        self.gBest_curve = []  # 每轮的最优解的适应度值
        self.gBest_X = np.zeros(dim)  # 每次迭代得到的最佳位置

    def BorderCheck_UpdateFitness(self):
        """搜索空间边界的检查与纠正与更新最优适应度值"""
        for i in range(self.population_size):
            for j in range(self.dim):
                if (self.X[i, j] > self.UB[j]):
                    # 不符合上边界
                    self.X[i, j] = self.UB[j]

                if (self.X[i, j] < self.LB[j]):
                    # 不符合下边界
                    self.X[i, j] = self.LB[j]
                # 计算所有物种的fitness
                fitness = self.CalculateFitness(self.X[i, :])  # 获取物种的适应度值
                # 更新最优解的适应度值和物种的最优解的位置信息
                if fitness < self.gBest_score:
                    self.gBest_score = fitness
                    self.gBest_X = self.X[i, :].copy()

    def Initialize_population(self):
        """initialize population 初始化种群数据"""
        # 初始化种群数据X
        for i in range(self.population_size):
            for j in range(self.dim):
                self.X[i, j] = (self.UB[j] - self.LB[j]) * \
                    np.random.random() + self.LB[j]

    def CalculateFitness(self, X):
        """自定义计算适应度函数,计算单个鲸鱼的适应度值"""
        fitness = self.FitFunction(X)
        return fitness

    def Optimize(self):
        self.Initialize_population()
        t = 1
        iter = 0
        while iter < self.max_iter:
            # 固定工作
            for i in range(self.population_size):
                for j in range(self.dim):
                    # 检查边界
                    for j in range(self.dim):
                        if (self.X[i, j] > self.UB[j]):
                            # 不符合上边界
                            self.X[i, j] = self.UB[j]

                        if (self.X[i, j] < self.LB[j]):
                            # 不符合下边界
                            self.X[i, j] = self.LB[j]
                    # 计算所有物种的fitness
                    fitness = self.CalculateFitness(self.X[i, :])  # 获取物种的适应度值
                    iter = iter+1
                    # 更新最优解的适应度值和物种的最优解的位置信息
                    if fitness < self.gBest_score:
                        self.gBest_score = fitness
                        self.gBest_X = self.X[i, :].copy()

            a = 2 * (1 - (iter / self.max_iter))   # 定义常数a,线性从2变为0
            a2 = -1 - (iter / self.max_iter)       # a2线性从-1~-2
            # 更新鲸鱼的位置信息
            for i in range(self.population_size):
                # 自定义三个过程的所需参数
                p = np.random.rand()
                r_1 = np.random.rand()
                r_2 = np.random.rand()
                A = 2 * a * r_1 - a
                C = 2 * r_2
                # l = 2 * np.random.rand() - 1  # 随机l数值[-1,1]
                l = (a2-1) * np.random.rand() + 1  # 官网给的版本
                for j in range(self.dim):
                    if (p < 0.5):
                        # 包围猎物与寻找猎物
                        if (abs(A) < 1):
                            # 缩小环绕更新位置  包围猎物
                            D_Best1 = abs(C * self.gBest_X[j] - self.X[i, j])
                            self.X[i, j] = self.gBest_X[j] - A * D_Best1
                        else:
                            # 寻找猎物,搜寻猎物
                            rand_index = np.random.randint(
                                low=0, high=self.population_size)
                            X_rand = self.X[rand_index, :]
                            D_Rand = abs(X_rand[j] * C - self.X[i, j])
                            self.X[i, j] = X_rand[j] - A * D_Rand
                    else:
                        # 螺旋式环绕更新位置
                        D_Best2 = abs(self.gBest_X[j]-self.X[i, j])
                        self.X[i, j] = D_Best2 * \
                            np.exp(self.b * l) * np.cos(2 *
                                                        np.pi * l) + self.gBest_X[j]

            # 更新每轮的最优解的适应度值
            self.gBest_curve.append(self.gBest_score)
            if (t % 10 == 0):
                print('第{}轮训练完成!, 该轮最优解的适应度值为: {}'.format(
                    t, self.gBest_curve[t-1]))
            t += 1

        return self.gBest_curve, self.gBest_X


if __name__ == '__main__':
    population_size = 30
    dim = 30
    max_iter = dim*10000
    LB = -10 * np.ones(dim)
    UB = 10 * np.ones(dim)
    Func = CEC(year=2017, dim=dim, fun_num=2)
    FitFunction = Func.fobj
    b = 2
    Gbest_curve, GbestX = WOA(FitFunction, LB=LB, UB=UB,
                              dim=dim, b=b, population_size=population_size, max_iter=max_iter).Optimize()
    print(Gbest_curve)
    # 可视化训练图像
    # FitnessPlot.Plot_2D('WOA', Gbest_curve)
